import cv2
import mediapipe as mp
import time
from tkinter import Tk, filedialog

# 初始化Mediapipe
mp_pose = mp.solutions.pose
mp_drawing = mp.solutions.drawing_utils
pose = mp_pose.Pose()

# 弹出文件选择对话框
root = Tk()
root.withdraw()
file_path = filedialog.askopenfilename(title="选择视频文件", filetypes=[("视频文件", "*.mp4")])

# 打开选择的视频文件
cap = cv2.VideoCapture(file_path)

# 初始化变量
pushup_count = 0
is_pushup_up = True

# 获取视频帧大小
frame_width = int(cap.get(3))
frame_height = int(cap.get(4))

# 倒计时
for i in range(3, 0, -1):
    ret, frame = cap.read()
    if not ret:
        continue

    # 显示带有动画效果的倒计时
    for size in range(10, 101, 10):
        countdown_text = str(i) if i > 1 else "1"
        font_size = size / 10.0
        font_thickness = size // 30 + 1

        font = cv2.FONT_HERSHEY_SIMPLEX
        text_size = cv2.getTextSize(countdown_text, font, font_size, font_thickness)[0]
        text_x = int((frame_width - text_size[0]) / 2)
        text_y = int((frame_height + text_size[1]) / 2)

        frame = cv2.rectangle(frame, (0, 0), (frame_width, frame_height), (255, 255, 255), -1)  # 添加白色背景
        frame_with_text = frame.copy()
        cv2.putText(frame_with_text, countdown_text, (text_x, text_y), font, font_size, (255, 0, 0), 12)
        cv2.imshow('count', frame_with_text)
        cv2.waitKey(50)

print("开始！")

while True:
    ret, frame = cap.read()
    if not ret:
        break  # 视频读取完毕，退出循环

    # 转换为RGB
    rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    # 进行姿势检测
    results = pose.process(rgb_frame)

    if results.pose_landmarks:
        # 获取关键点
        landmarks = results.pose_landmarks.landmark

        # 计算肩膀与地面的距离
        shoulder_y = int(landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER].y * frame_height)
        hip_y = int(landmarks[mp_pose.PoseLandmark.LEFT_HIP].y * frame_height)
        distance = hip_y - shoulder_y

        # 检测是否处于俯卧撑上升或下降阶段
        if is_pushup_up and distance > 50:
            is_pushup_up = False
        elif not is_pushup_up and distance < 20:
            is_pushup_up = True
            pushup_count += 1

        # 获取文本的大小
        text_size = cv2.getTextSize(f'count: {pushup_count}', cv2.FONT_HERSHEY_SIMPLEX, 1, 2)[0]

        # 计算矩形框的位置和大小
        rect_position = (10, 30 - text_size[1])
        rect_size = (text_size[0] + 10, text_size[1] + 10)  # 可根据需要调整填充大小

        # 绘制矩形框
        cv2.rectangle(frame, rect_position, (rect_position[0] + rect_size[0], rect_position[1] + rect_size[1]), (0, 255, 0), -1)

        # 在图像上显示俯卧撑计数
        cv2.putText(frame, f'count: {pushup_count}', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2)

        # 在图像上绘制关键点
        mp_drawing.draw_landmarks(frame, results.pose_landmarks, mp_pose.POSE_CONNECTIONS)

    # 显示图像
    cv2.imshow('count', frame)

    # 按下空格键退出
    if cv2.waitKey(1) & 0xFF == ord(' '):
        break

# 释放视频文件资源
cap.release()
cv2.destroyAllWindows()
